#本脚本的目的是为了将seg_image_transform_json转化出的voc数据格式转化为labelme可以读写的格式

import xml.etree.ElementTree as ET
import json
import os
from pathlib import Path
from tqdm import tqdm

def voc_segmentation_to_labelme(voc_xml_file, output_dir):
    tree = ET.parse(voc_xml_file)
    root = tree.getroot()

    filename = root.find('filename').text
    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)

    labelme_data = {
        "version": "5.3.1",
        "flags": {},
        "shapes": [],
        "imagePath": filename,
        "imageData": None,
        "imageHeight": height,
        "imageWidth": width
    }

    for obj in root.findall('object'):
        label = obj.find('name').text
        # 假设多边形点存储在segm/polygon节点下
        polygon = obj.find('polygon')
        if polygon is not None:
            points = []
            coords = list(polygon)  # 获取所有子节点
            for i in range(0, len(coords), 2):  # 步长为2以匹配x,y对
                x = int(coords[i].text)
                y = int(coords[i+1].text)
                points.append([x, y])

            shape = {
                "label": label,
                "points": points,
                "group_id": None,
                "shape_type": "polygon",
                "flags": {}
            }

            labelme_data['shapes'].append(shape)
    json_filename = os.path.join(output_dir, os.path.splitext(os.path.basename(voc_xml_file))[0] + '.json')
    with open(json_filename, 'w') as json_file:
        json.dump(labelme_data, json_file, indent=2)


if __name__ == "__main__":
    voc_xml_dir = '/home/luoluoluo/data/dataset/caixukun_mask'
    output_dir = '/home/luoluoluo/data/dataset/caixukun'
    for voc_xml_file in tqdm(Path(voc_xml_dir).glob("*.xml")):
        voc_segmentation_to_labelme(str(voc_xml_file), output_dir)
